package main

import (
	"context"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/joho/godotenv"
	openai "github.com/sashabaranov/go-openai"
)

var (
	client    *openai.Client
	chatModel string = "gpt-3.5-turbo"
)

func init() {
	godotenv.Load(".env")
	config := openai.DefaultConfig(os.Getenv("OPENAI_API_KEY"))
	if os.Getenv("OPENAI_BASE_URL") != "" {
		config.BaseURL = strings.TrimRight(os.Getenv("OPENAI_BASE_URL"), "/")
	}
	client = openai.NewClientWithConfig(config)
	if os.Getenv("OPENAI_CHAT_MODEL") != "" {
		chatModel = os.Getenv("OPENAI_CHAT_MODEL")
	}
	log.Printf("config loaded: url=%s; model=%s\n", config.BaseURL, chatModel)
}

// 定义升级器，从HTTP协议升级到WebSocket协议
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许跨域
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// 定义WebSocket处理器
func wsHandler(w http.ResponseWriter, r *http.Request) {
	// 升级连接到WebSocket协议
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Upgrade error:", err)
		return
	}
	defer conn.Close()

	for {
		// 读取新的消息
		_, p, err := conn.ReadMessage()
		if err != nil {
			log.Println("Read error:", err)
			break
		}

		// 处理接收到的query
		log.Printf("Received: %s", p)

		ctx := context.Background()
		prompt := string(p)
		stream, err := client.CreateChatCompletionStream(ctx, openai.ChatCompletionRequest{
			Model: chatModel,
			Messages: []openai.ChatCompletionMessage{{
				Role:    openai.ChatMessageRoleUser,
				Content: prompt,
			}},
			Stream: true,
		})
		if err != nil {
			log.Println("OpenAI error:", err)
			continue
		}

		log.Println("stream start...")
		for {
			response, err := stream.Recv()
			if err != nil {
				log.Println("Stream error:", err)
				break
			}

			answer := response.Choices[0].Delta.Content
			if err := conn.WriteMessage(websocket.TextMessage, []byte(answer)); err != nil {
				log.Println("Write error:", err)
				break
			}
		}
		if err := conn.WriteMessage(websocket.TextMessage, []byte{0}); err != nil {
			log.Println("Write error:", err)
			break
		}
		log.Println("stream end...")
	}
}

// 定义静态文件处理器
func staticHandler(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	if strings.HasSuffix(path, "/") {
		path = path + "index.html"
	}
	fp := filepath.Join("client", path)

	if _, err := filepath.Abs(fp); err != nil {
		http.NotFound(w, r)
		return
	}
	http.ServeFile(w, r, fp)
}

func main() {
	http.HandleFunc("/", staticHandler)
	http.HandleFunc("/ws", wsHandler)

	log.Fatal(http.ListenAndServe(":8080", nil))
}
